// Copyright (c) 2019 Presto Labs Pte. Ltd.
// Author: jaewon

#pragma once

#include <algorithm>
#include <deque>
#include <optional>
#include <utility>

#include "coin2/base/log.h"

namespace coin2::base::math {

class TimeWeightedMovingAverage {
 public:
  using TimeValuePairType = std::pair<int64_t, double>;
  using Container = std::deque<TimeValuePairType>;
  using Iterator = typename Container::iterator;
  using ConstIterator = typename Container::const_iterator;

  TimeWeightedMovingAverage() = default;

  explicit TimeWeightedMovingAverage(int64_t window_size) : window_size_(window_size) {
    CHECK_GT_THROW(window_size, 0);
  }

  ~TimeWeightedMovingAverage() = default;

  void Push(int64_t time, double value) {
    if (!window_.empty()) {
      const auto& last_pushed = window_.back();
      sum_ -= head_tv_;
      sum_ += static_cast<double>(time - last_pushed.first) * last_pushed.second;
    }

    head_tv_ = 0.;
    window_.push_back(std::make_pair(time, value));
    UpdateTime(time);
  }

  void PushTime(int64_t time) {
    if (!window_.empty()) {
      const auto& last_pushed = window_.back();
      sum_ -= head_tv_;
      head_tv_ = static_cast<double>(time - last_pushed.first) * last_pushed.second;
      sum_ += head_tv_;
    }

    UpdateTime(time);
  }

  void ResizeWindow(int64_t window_size) {
    CHECK_GT_THROW(window_size, 0);
    window_size_ = window_size;
    if (!window_.empty()) {
      PushTime(window_.back().first);
    }
  }

  void Clear() {
    window_.clear();
    sum_ = 0.;
    tail_tv_ = 0.;
    head_tv_ = 0.;
    num_popped_ = 0;
    last_popped_ = std::nullopt;
  }

  double tw_sum() const { return sum_; }
  double tw_mean() const { return sum_ / static_cast<double>(window_size_); }

  size_t num_popped() const { return num_popped_; }
  const std::optional<TimeValuePairType>& last_popped() const { return last_popped_; }

  int64_t window_size() const { return window_size_; }

  bool empty() const { return window_.empty(); }
  size_t size() const { return window_.size(); }

  const TimeValuePairType& front() const { return window_.front(); }
  const TimeValuePairType& back() const { return window_.back(); }

  Iterator begin() { return window_.begin(); }
  Iterator end() { return window_.end(); }
  ConstIterator begin() const { return window_.begin(); }
  ConstIterator end() const { return window_.end(); }

 private:
  void UpdateTime(int64_t time) {
    if (num_popped_ > 0) {
      sum_ -= tail_tv_;
    }

    bool empty = true;
    for (auto it = window_.begin(); it != window_.end(); ++it) {
      if (time - it->first < window_size_) {
        empty = false;
        if (it != window_.begin()) {
          window_.erase(window_.begin(), it);
        }
        break;
      } else {
        ++num_popped_;
        last_popped_ = *it;

        auto next = it;
        ++next;
        if (next != window_.end()) {
          sum_ -= static_cast<double>(next->first - it->first) * last_popped_.value().second;
        } else {
          sum_ -= head_tv_;
          head_tv_ = 0.;
        }
      }
    }
    if (empty) {
      window_.clear();
      sum_ = 0.;
      head_tv_ = 0.;
      tail_tv_ = 0.;
    }

    if (num_popped_ > 0) {
      int64_t tail_time = time;
      if (!window_.empty()) {
        tail_time = window_.front().first;
      }
      double weight = std::min(window_size_, window_size_ - (time - tail_time));
      tail_tv_ = weight * last_popped_.value().second;
      sum_ += tail_tv_;
    }
  }

  int64_t window_size_ = 0;
  size_t num_popped_ = 0;
  std::optional<TimeValuePairType> last_popped_;

  double sum_ = 0.;
  double head_tv_ = 0.;
  double tail_tv_ = 0.;
  Container window_;
};

}  // namespace coin2::base::math
